
import logging
from transformers import (
    TrainerCallback,
    TrainingArguments,
    TrainerState,
    TrainerControl
)

logger = logging.getLogger(__name__)
class LoggingCallback(TrainerCallback):
    def on_step_end(self,args: TrainingArguments, state: TrainerState,control:TrainerControl,**kwargs):
        if state.global_step % args.logging_steps == 0:
            logger.info(f"Step {state.global_step}: Loss = {state.log_history[-1].get('loss',None)},Learning Rate = {state.log_history[-1].get('learn_rate',None)}")
def get_callbacks(training_args,model_args,script_args):
    callbacks = [LoggingCallback()]        
    return callbacks
